import os
import tarfile
import argparse
import time

from taichu_storage import command
import docker
import json


def upload_to_obs(src, dest):
    command.upload_to_obs(src, dest)


def parse_docker_config():
    with open("/etc/docker/config.json") as f:
        data = json.load(f)

    for registry, auth_info in data.get('auths', {}).items():
        username = auth_info.get('username')
        password = auth_info.get('password')
        break

    return registry, username, password


def pack(arg):
    os.chdir(arg.workdir)
    print(f'change to workdir: {os.path.abspath(arg.workdir)}')

    output_dir = os.path.dirname(arg.output)
    if not output_dir.startswith("obs://"):
        os.makedirs(output_dir, exist_ok=True)

    if len(arg.add_image) > 0:
        client = docker.DockerClient(base_url=f'tcp://docker-daemon.infra:2375', timeout=2 * 60 * 60)
        registry, username, password = parse_docker_config()
        client.login(registry=registry, username=username, password=password)
        print("docker login success")
        for i in arg.add_image:
            print(f"pulling image {i}")
            image = client.images.pull(i)
            with open('model-service.tar', 'wb') as f:
                print("writing model-service.tar")
                for chunk in image.save():
                    f.write(chunk)

    output_file_name = os.path.basename(arg.output)
    with tarfile.open(output_file_name, 'w') as tf:
        if os.path.isfile('model-service.tar'):
            print(f"adding model-service.tar")
            tf.add('model-service.tar', 'model-service.tar')

        for d in arg.add_dir:
            print(f"adding dir {d}")
            d = d.strip()
            if d.startswith("obs://"):
                local_dirname = os.path.basename(d)
                if local_dirname is None or local_dirname == "":
                    local_dirname = os.path.basename(d[:-1])
                command.download_from_obs(d, f"./{local_dirname}/")
                d = f"./{local_dirname}"

            if os.path.isdir(d):
                tf.add(d, os.path.basename(d), recursive=True)
            print(f"adding dir {d} success")

        for f in arg.add_file:
            print(f"adding file {f}")
            local_file_name = f.strip()
            if local_file_name.startswith("obs://"):
                command.download_from_obs(local_file_name, f"./{os.path.basename(local_file_name)}")
                local_file_name = os.path.basename(local_file_name)

            tf.add(local_file_name, f"./{local_file_name}")
            print(f"add file {f} success")

    if arg.output.startswith("obs://"):
        print(f"uploading {output_file_name} to {arg.output}")
        upload_to_obs(output_file_name, arg.output)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--workdir', '-w', type=str, default='.', help="workdir")
    parser.add_argument('--add_dir', type=str, action='append', default=[], required=False, help="add path")
    parser.add_argument('--output', '-o', type=str, required=False, default='out.tar')
    parser.add_argument('--add_file', type=str, action='append', default=[], help="add file")
    parser.add_argument('--add_image', type=str, action='append', default=[], help="add image")

    args = parser.parse_args()

    pack(args)
    print('done!')
